package com.qupeng.demo.kafka.kafkaspringbootproducer.admin;

import com.qupeng.demo.kafka.kafkaspringbootproducer.admin.model.Topic;
import com.qupeng.demo.kafka.kafkaspringbootproducer.config.KafkaProperties;
import com.qupeng.demo.kafka.kafkaspringbootproducer.model.CommonResponse;
import com.qupeng.demo.kafka.kafkaspringbootproducer.admin.model.Partition;
import org.apache.kafka.clients.admin.*;
import org.apache.kafka.common.ElectionType;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.kafka.core.KafkaAdmin;
import org.springframework.web.bind.annotation.*;

import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/admin/sync")
public class PartitionSyncAdminController {

    private static final Logger LOGGER = LoggerFactory.getLogger(PartitionSyncAdminController.class);

    @Autowired
    private KafkaAdmin kafkaAdmin;

    @Autowired
    private AdminClient adminClient;

    @Autowired
    private KafkaProperties kafkaProperties;

    @PostMapping("/topics/{name}/partitions")
    public CommonResponse<List<Partition>> topics(@PathVariable String name) {
        int partitionCount = kafkaAdmin.describeTopics(name).get(name).partitions().size();
        Map<String, NewPartitions> newPartitionsMap = new HashMap<>();
        newPartitionsMap.put(name, NewPartitions.increaseTo(++partitionCount));

        try {
            adminClient.createPartitions(newPartitionsMap).all().get();
            List<Partition> partitions = kafkaAdmin.describeTopics(name).get(name).partitions().stream().map(partitionInfo -> {
                return Util.convert(partitionInfo);
            }).collect(Collectors.toList());
            return new CommonResponse<>("kafka-admin", partitions);
        } catch (InterruptedException | ExecutionException  e) {
            return new CommonResponse<>("001-partition-management", e.getMessage(), "kafka-admin");
        }
    }

    @PutMapping("/topics/{name}/partitions")
    public CommonResponse<Topic> electLeader(@PathVariable String name, @RequestBody Map<String, String> body) {
        if (body.get("operation").equals("electLeaders")) {
            Set<TopicPartition> electableTopics = new HashSet<>();
            electableTopics.add(new TopicPartition(name, 0));
            try {
                adminClient.electLeaders(ElectionType.PREFERRED, electableTopics).all().get();
            } catch (InterruptedException | ExecutionException  e) {
                return new CommonResponse<>("002-partition-management", e.getMessage(), "kafka-admin");
            }
        } else if (body.get("operation").equals("reassignPartitions")) {
            Map<TopicPartition, Optional<NewPartitionReassignment>> reassignment = new HashMap<>();
            reassignment.put(new TopicPartition(name, 0), Optional.of(new NewPartitionReassignment(Arrays.asList(0, 1))));
            try {
                adminClient.alterPartitionReassignments(reassignment).all().get();
            } catch (InterruptedException | ExecutionException  e) {
                return new CommonResponse<>("003-partition-management", e.getMessage(), "kafka-admin");
            }
        }

        try {
            Map<TopicPartition, PartitionReassignment> result = adminClient.listPartitionReassignments().reassignments().get();
            TopicDescription topicDescription = adminClient.describeTopics(Collections.singletonList(name)).topicNameValues().get(name).get();
            Topic topic = new Topic();
            BeanUtils.copyProperties(topicDescription, topic);
            for (TopicPartitionInfo partitionInfo : topicDescription.partitions()) {
                Partition partition = new Partition();
                BeanUtils.copyProperties(partitionInfo, partition);
                topic.getPartitions().add(partition);
            }
            return new CommonResponse<>("kafka-admin", topic);
        } catch (InterruptedException | ExecutionException  e) {
            return new CommonResponse<>("004-partition-management", e.getMessage(), "kafka-admin");
        }

    }
}
